import json
import os
from pathlib import Path


INPUT_SCHEMES = [
    # 'ecg',
    # 'cwt',
    # 'ecg+cwt',
    # 'ecg+rr',
    # 'cwt+rr',
    'ecg+cwt+rr'
]
POSITIONAL_ENCODING_SCHEMES = [
	# 'zero_addition_learned',
	'timestamp_concat_after',
	'linspace_concat_after',
	'standard_addition',
]
DATASETS = [
	# 'incart_stratified_standard_original',
	'incart_stratified_standard_train_100hz_val_500hz',
	# 'incart_stratified_standard_25hz_5sr'
]

MODEL_TYPE = 'transformer'
LATE_FUSION = True

DATA_ROOT = Path('/project_ghent/Timo/digihealth_asia/datasets')

if __name__ == '__main__':
    job_scripts_dir = Path(os.path.dirname(os.path.realpath(__file__)))
    job_template_path = job_scripts_dir / 'job_template.txt'
    job_dict = json.load(open(job_template_path, 'r'))
    command_str = job_dict['request']['docker']['command']

    if MODEL_TYPE == 'cnn':
        POSITIONAL_ENCODING_SCHEMES = [None,]
    
    for input_scheme in INPUT_SCHEMES:
        for positional_encoding_scheme in POSITIONAL_ENCODING_SCHEMES:
            for dataset in DATASETS:
                data_path = DATA_ROOT / dataset
                if positional_encoding_scheme is not None:
                    job_name = f'{input_scheme}_{positional_encoding_scheme}_{dataset}_{MODEL_TYPE}'
                else:
                    job_name = f'{input_scheme}_{dataset}_{MODEL_TYPE}'
                if 'rr' in input_scheme and LATE_FUSION:
                    job_name += '_late_fusion'
                current_command_str = command_str + f' --input_type {input_scheme}'
                current_command_str += f' --model {MODEL_TYPE}'
                if positional_encoding_scheme is not None:
                    current_command_str += f' --positional_encoding_scheme {positional_encoding_scheme}'
                current_command_str += f' --late_fusion {LATE_FUSION}'
                current_command_str += f' --data_path {str(data_path)}'
                current_job_dict = job_dict.copy()
                current_job_dict['request']['docker']['command'] = current_command_str
                current_job_dict['name'] = f'heartbeat_classification_{job_name}'
                job_path = job_scripts_dir / f'{job_name}.json'
                with open(job_path, 'w') as f:  
                    json.dump(current_job_dict, f, indent=4)
            
